use std::path::{Path, PathBuf};
use std::sync::Arc;

use cudarc::driver::{
    CudaContext, CudaModule, CudaSlice, CudaStream, DriverError, LaunchConfig, PushKernelArg,
};

use glium::Display;
use glium::winit::event_loop::EventLoop;
use glutin::surface::WindowSurface;

use use_cudarc::framebuffer_renderer::{CudaRendererOps, FrameBufferRenderer};
use use_cudarc::model::{GpuModel, Model};
use use_cudarc::tgaimage::{ColorFormat, TGAImage};
use use_cudarc::window_app::App;

struct ModelSkeletonRenderer {
    frame_buffer_size: (usize, usize),
    frame_buffer: Vec<u8>,
    frame_buffer_gpu: CudaSlice<u8>,

    cuda_stream: Arc<CudaStream>,
    module: Arc<CudaModule>,
}

impl ModelSkeletonRenderer {
    fn new<P: AsRef<Path>>(
        frame_buffer_size: (usize, usize),
        shader: P,
    ) -> Result<Self, DriverError> {
        let cuda_ctx = CudaContext::new(0)?;
        let cuda_stream = cuda_ctx.default_stream();

        let frame_buffer = vec![[0, 0, 0, 255]; frame_buffer_size.0 * frame_buffer_size.1].concat();
        let frame_buffer_gpu = cuda_stream.memcpy_stod(&frame_buffer)?;

        let ptx = cudarc::nvrtc::compile_ptx(std::fs::read_to_string(shader).unwrap()).unwrap();
        let module = cuda_ctx.load_module(ptx)?;

        Ok(Self {
            frame_buffer_size,
            frame_buffer,
            frame_buffer_gpu,
            cuda_stream,
            module,
        })
    }

    fn render_model(&mut self, model: &GpuModel) -> Result<(), DriverError> {
        let kernel = self.module.load_function("render_model_kernel")?;

        let face_num = (model.facet_vrt.len() / 3) as u32;
        let (width, height) = (
            self.frame_buffer_size.0 as u32,
            self.frame_buffer_size.1 as u32,
        );

        // Launch the kernel
        let mut builder = self.cuda_stream.launch_builder(&kernel);
        builder.arg(&model.verts);
        builder.arg(&model.facet_vrt);
        builder.arg(&mut self.frame_buffer_gpu);
        builder.arg(&face_num);
        builder.arg(&width);
        builder.arg(&height);
        unsafe { builder.launch(LaunchConfig::for_num_elems(face_num)) }?;

        self.cuda_stream
            .memcpy_dtoh(&self.frame_buffer_gpu, &mut self.frame_buffer)
    }
}

struct Renderer {
    skeleton_renderer: ModelSkeletonRenderer,
    models: Vec<GpuModel>,
}

impl CudaRendererOps for Renderer {
    fn gen_framebuffer(&mut self, frame_index: usize, fb_size: (usize, usize)) -> Vec<u8> {
        assert!(self.skeleton_renderer.frame_buffer_size == fb_size);

        for model in &self.models {
            self.skeleton_renderer.render_model(model).unwrap()
        }

        let pixels = self.skeleton_renderer.frame_buffer.clone();
        let (width, height) = self.skeleton_renderer.frame_buffer_size;

        if frame_index == 0 || frame_index == 100 {
            TGAImage::new(width, height, ColorFormat::RGBA)
                .set_data(&pixels)
                .write_tga_file(format!("frame{frame_index}.tga"), true)
                .unwrap()
        }

        pixels
    }
}

fn main() {
    let args: Vec<String> = std::env::args().collect();

    let event_loop = EventLoop::new().unwrap();
    let bf_size = (2400, 2400);

    let skeleton_renderer = ModelSkeletonRenderer::new(
        bf_size,
        PathBuf::from(file!()).parent().unwrap().join("kernel.cu"),
    )
    .unwrap();

    let mut models = vec![];

    let model_dir = &args[1];
    assert!(std::fs::exists(model_dir).unwrap());
    for entry in std::fs::read_dir(model_dir).unwrap() {
        let entry = entry.unwrap();
        let path = entry.path();

        if path.is_file() && path.extension().map(|ext| ext == "obj").unwrap_or(false) {
            models.push(
                GpuModel::from_model(
                    &skeleton_renderer.cuda_stream,
                    &Model::from_file(&path).unwrap(),
                )
                .unwrap(),
            );
        }
    }

    let mut app = App::new(
        Box::new(
            move |display: &Display<WindowSurface>| -> FrameBufferRenderer<Renderer> {
                FrameBufferRenderer::<Renderer>::new(
                    display,
                    bf_size,
                    Renderer {
                        skeleton_renderer,
                        models,
                    },
                )
            },
        ),
        "SKELETON RENDER",
        bf_size,
    );
    event_loop.run_app(&mut app).unwrap()
}
